import torch
import torchvision

# config data
batch_size = 512

# step1:load data
train_loader = torch.utils.data.DataLoader(
    torchvision.datasets.MNIST('mnist_data', train=True, download=True,
                               transform=torchvision.transforms.Compose([
                                   torchvision.transforms.ToTensor(),
                                   torchvision.transforms.Normalize(
                                       (0.1307,), (0.3081,))
                               ])),
    batch_size=batch_size, shuffle=True)

test_loader = torch.utils.data.DataLoader(
    torchvision.datasets.MNIST('mnist_data/', train=False, download=True,
                               transform=torchvision.transforms.Compose([
                                   torchvision.transforms.ToTensor(),
                                   torchvision.transforms.Normalize(
                                       (0.1307,), (0.3081,))
                               ])),
    batch_size=batch_size, shuffle=True)

# x, y = next(iter(train_loader))
# print(x.shape, y.shape)
# plot_image(x, y, 'image sample')

# step2:build model


class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # model
        self.fc1 = torch.nn.Linear(28*28, 256)
        self.fc2 = torch.nn.Linear(256, 64)
        self.fc3 = torch.nn.Linear(64, 10)

    def forward(self, x):
        # h1=relu(xw1+b1)
        x = torch.nn.functional.relu(self.fc1(x))
        # h2=relu(h1w2+b2)
        x = torch.nn.functional.relu(self.fc2(x))
        # h3=relu(h2w3+b3)
        x = self.fc3(x)
        return x


def one_hot(label, depth=10):
    '''
    '''
    out = torch.zeros(label.size(0), depth)
    idx = torch.LongTensor(label).view(-1, 1)
    out.scatter_(dim=1, index=idx, value=1)
    return out


# train
# model
net = Net()
# Optimizer
optimizer = torch.optim.SGD(net.parameters(), lr=0.01, momentum=0.9)

for epoch in range(3):
    for batch_idx, (x, y) in enumerate(train_loader):
        print(x.size())
        
        x = x.view(x.size(0), 28*28)
        print(x.size())
        break
        out = net(x)
        # [b,10]
        y_onehot = one_hot(y)
        # loss = mse(out,y_onehot)
        loss = torch.nn.functional.mse_loss(out, y_onehot)
        optimizer.zero_grad()
        loss.backward()
        # w' = w - lr * grad
        optimizer.step()
        if batch_idx % 10 == 0:
            print(epoch, batch_idx, loss.item())

# get optimal [w1 w2 w3 b1 b2 b3]
total_correct = 0
for x, y in test_loader:
    x = x.view(x.size(0), 28*28)
    out = net(x)
    pred = out.argmax(dim=1)
    correct = pred.eq(y).sum().float()
    total_correct += correct

total_sum = len(test_loader.dataset)
acc = total_correct / total_sum
print('acc=', acc)
